#define DTYPE_X half
#define DTYPE_Y half
#define DTYPE_Z half
#include "whole_reduce_sum_tiling.h"
#include "kernel_operator.h"
constexpr int32_t BUFFER_NUM = 2;

class KernelWholeReduceSum {
    
    
    public:
        __aicore__ inline KernelWholeReduceSum() {}
        __aicore__ inline void Init(GM_ADDR x, GM_ADDR z, uint32_t coreInputDataNum, uint32_t coreOutputDataNum) {

            if (0 == AscendC::GetBlockIdx()) {
                AscendC::printf("Init start: coreInputDataNum=%u, coreOutputDataNum=%u\n", coreInputDataNum, coreOutputDataNum);
                AscendC::printf("Init start: block number=%u, blockIdx=%u\n", AscendC::GetBlockNum(), AscendC::GetBlockIdx());

            }


            uint32_t coreNum = AscendC::GetBlockIdx();
            uint32_t xGmIndex = coreInputDataNum * AscendC::GetBlockIdx();
            
            uint32_t zGmIndex = coreOutputDataNum * AscendC::GetBlockIdx();


            this->coreInputDataNum = coreInputDataNum;
            this->coreOutputDataNum = coreOutputDataNum;

            
            xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + xGmIndex, this->coreInputDataNum);
            
            zGm.SetGlobalBuffer((__gm__ DTYPE_Z*)z + zGmIndex, this->coreOutputDataNum);

            pipe.InitBuffer(inQueueX, BUFFER_NUM, this->coreInputDataNum * sizeof(DTYPE_X));
            
            pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->coreOutputDataNum * sizeof(DTYPE_Z));
        }

        __aicore__ inline void Process() {
            int32_t loopCount = 1;
            
            for (int32_t i = 0; i < loopCount; i++) {
                CopyIn(i);
                Compute(i);
                CopyOut(i);
            }
        }

    private:
        __aicore__ inline void CopyIn(int32_t progress) {
            AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
            
            AscendC::DataCopy(xLocal, xGm[progress * this->coreInputDataNum], this->coreInputDataNum);

            
            inQueueX.EnQue(xLocal);
            
        }

        __aicore__ inline void Compute(int32_t progress) {
            AscendC::LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
            
            AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.AllocTensor<DTYPE_Z>();
            AscendC::WholeReduceSum<DTYPE_X>(zLocal, xLocal, 128, 16, 1, 1, 8);
            outQueueZ.EnQue<DTYPE_Z>(zLocal);
            inQueueX.FreeTensor(xLocal);
            

        }

        __aicore__ inline void CopyOut(int32_t progress) {
            AscendC::LocalTensor<DTYPE_Z> zLocal = outQueueZ.DeQue<DTYPE_Z>();
            AscendC::DataCopy(zGm[progress * this->coreOutputDataNum], zLocal, this->coreOutputDataNum);
            outQueueZ.FreeTensor(zLocal);
        }

    private:
        AscendC::TPipe pipe;
        AscendC::TQue<AscendC::QuePosition::VECIN, BUFFER_NUM> inQueueX;
        AscendC::TQue<AscendC::QuePosition::VECOUT, BUFFER_NUM> outQueueZ;
        AscendC::GlobalTensor<DTYPE_X> xGm;
        
        AscendC::GlobalTensor<DTYPE_Z> zGm;
        uint32_t coreInputDataNum;
        uint32_t coreOutputDataNum;

};

extern "C" __global__ __aicore__ void whole_reduce_sum(GM_ADDR x, GM_ADDR z, WholeReduceSumTilingData tiling) {
    KernelWholeReduceSum op;
    op.Init(x, z, tiling.coreInputDataNum, tiling.coreOutputDataNum);

    op.Process();
}
